# 三维Laplacian方程松弛法求解
# \laplacian u = b
# 参考 https://zhuanlan.zhihu.com/p/144331014
# 可能有bug
# Gitee Repo

import numpy as np
import torch
import scipy

import plotly.graph_objects as go
from plotly.offline import plot

import time

torch.set_default_device('cuda' if torch.cuda.is_available() else 'cpu')

time0 = time.time()

L = 2;
dx = 0.1;

x,y,z = torch.meshgrid(torch.arange(-L,L,dx),torch.arange(-L,L,dx),torch.arange(-L,L,dx),indexing='ij');
n = int(x.shape[0]);

u0 = torch.zeros((n,n,n))
u0[0,:,:]=0
u0[n-1,:,:]=0
u0[:,0,:]=0
u0[:,n-1,:]=0
u0[:,:,0]=0
u0[:,:,n-1]=0

# 边界条件
u0[:,:,n-1] = torch.cos(torch.pi/2/L*x[:,:,n-1])*torch.cos(torch.pi/2/L*y[:,:,n-1])

# 右侧向量
b = torch.zeros((n,n,n))
# b[(x+0.5*L)**2+y**2+z**2<(0.2*L)**2] = 1;
# b[(x-0.5*L)**2+y**2+z**2<(0.2*L)**2] = -1;

u1 = u0.clone()

for tick in range(1000):
    u1[1:n-1,1:n-1,1:n-1] = (
    +u0[2:n,1:n-1,1:n-1] \
    +u0[0:n-2,1:n-1,1:n-1] \
    +u0[1:n-1,2:n,1:n-1] \
    +u0[1:n-1,0:n-2,1:n-1] \
    +u0[1:n-1,1:n-1,2:n] \
    +u0[1:n-1,1:n-1,0:n-2] \
    )/6 \
    - b[1:n-1,1:n-1,1:n-1] * (dx)**2
    
    err = torch.max(torch.abs(u0-u1))
    if err < 1e-4:
        break;
    if tick % 10 == 0:
        print(tick,err.item())
    u0,u1 = u1,u0;

time1 = time.time()
print(time1-time0)

def draw(x,y,z,u):
    import plotly.graph_objects as go
    from plotly.offline import plot

    downsample_factor = 3 #降低精度以避免过大的数据

    x_cpu = x.cpu().numpy()[::downsample_factor, ::downsample_factor, ::downsample_factor]
    y_cpu = y.cpu().numpy()[::downsample_factor, ::downsample_factor, ::downsample_factor]
    z_cpu = z.cpu().numpy()[::downsample_factor, ::downsample_factor, ::downsample_factor]
    u_cpu = u.cpu().numpy()[::downsample_factor, ::downsample_factor, ::downsample_factor]

    fig = go.Figure(data=go.Volume(
        x=x_cpu.flatten(),
        y=y_cpu.flatten(),
        z=z_cpu.flatten(),
        value=u_cpu.flatten(),
        opacity=0.1, # needs to be small to see through all surfaces
        surface_count=20, # needs to be a large number for good volume rendering
        ))

    plot(fig)
    scipy.io.savemat('res.mat',{'x':x_cpu,'y':y_cpu,'z':z_cpu,'u':u_cpu})
    return

draw(x,y,z,u1)
